from typing import Any, Callable, Iterable, Optional

from PyBirdViewCode import universal_ast_nodes as uast_nodes
from PyBirdViewCode.universal_ast.models import NotImplementedItem

from ..asty import nodes as asty_nodes
from ..asty.nodes import NodeType

# with open(r"C:\Users\houzh\Developing\opensource-projects\asty\output.json") as f:
#     tree: FileNode = FileNode.parse_raw(f.read())
#     tree


class GolangASTConverter:
    def __init__(
        self,
    ) -> None:
        # self.source_location_filter: Optional[Callable[[SourceLocation], bool]] = None
        self._handlers_map: dict[NodeType, Callable[[asty_nodes.BaseNode], Any]] = {
            NodeType.FILE: self._handle_file,
            NodeType.FUNC_DECL: self._handle_func_decl,
            NodeType.BLOCK_STMT: self._handle_block_stmt,
            NodeType.DECL_STMT: self._handle_decl_stmt,
            NodeType.GEN_DECL: self._handle_gen_decl,
            # NodeType.VALUE_SPEC: self._handle_value_spec,
            NodeType.ASSIGN_STMT: self._handle_assign_stmt,
            NodeType.EXPR_STMT: self._handle_expr_stmt,
            NodeType.IF_STMT: self._handle_if_stmt,
            NodeType.SWITCH_STMT: self._handle_switch_stmt,
            NodeType.CASE_CLAUSE: self._handle_case_clause,
            NodeType.RETURN_STMT: self._handle_return_stmt,
            NodeType.CALL_EXPR: self._handle_call_expr,
            NodeType.UNARY_EXPR: self._handle_unary_expr,
            NodeType.BINARY_EXPR: self._handle_binary_expr,
            NodeType.SELECTOR_EXPR: self._handle_selector_expr,
            NodeType.KEY_VALUE_EXPR: self._handle_key_value_expr,
            NodeType.IDENT: self._handle_ident,
            # NodeType.FIELD: self._handle_field,
            NodeType.BASIC_LIT: self._handle_basic_lit,
            NodeType.FUNC_LIT: self._handle_func_lit,
            NodeType.COMPOSITE_LIT: self._handle_composite_lit,
            NodeType.FUNC_TYPE: self._handle_func_type,
            NodeType.ARRAY_TYPE: self._handle_array_type,
            NodeType.STRUCT_TYPE: self._handle_structure_type,
            # NodeType.GEN_DECL: self._handle_gen_dec,
        }

        # Return value of execution
        self._ret_value = None
        # self._labels: Dict[str, LabelDesc] = {}

    def _handle_notimplemented(self, c: asty_nodes.BaseNode) -> NotImplementedItem:
        return NotImplementedItem(str(c.kind))

    def _handle_return_stmt(
        self, node: asty_nodes.ReturnStmtNode
    ) -> uast_nodes.ReturnStmt:
        return uast_nodes.ReturnStmt(
            self.eval_node_list(node.results) if node.results is not None else []
        )

    def _parse_field_list(
        self, node: asty_nodes.FieldListNode
    ) -> list[uast_nodes.ParamDecl]:

        if node.field_list is not None:
            decls = []
            field: asty_nodes.FieldNode
            for field in node.field_list:
                for name in field.names:
                    decls.append(
                        uast_nodes.ParamDecl(
                            self.eval_single_node(name),
                            self.eval_single_node(field.type),
                        )
                    )
            return decls
        else:
            return []

    def _parse_field(self, node: asty_nodes.FieldNode) -> list[uast_nodes.FieldDecl]:
        # for name in node.names:
        if node.names is None:
            return []
        return [
            uast_nodes.FieldDecl(
                self.eval_single_node(name),
                self.eval_single_node(node.type),
            )
            for name in node.names
        ]

    def _handle_structure_type(
        self, node: asty_nodes.StructTypeNode
    ) -> uast_nodes.StructDecl:
        fields = []
        if (node.fields is None) or (node.fields.field_list is None):
            pass
        else:
            for field in node.fields.field_list:
                fields.append(self._parse_field(field))
        return uast_nodes.StructDecl(name=None, fields=fields)

    def _handle_array_type(
        self, node: asty_nodes.ArrayTypeNode
    ) -> uast_nodes.ArrayType:
        return uast_nodes.ArrayType(
            self.eval_single_node(node.elt) if node.elt is not None else None,
            self.eval_single_node(node.len) if node.len is not None else None,
        )

    def _handle_func_type(self, node: asty_nodes.FuncTypeNode) -> uast_nodes.MethodType:
        return uast_nodes.MethodType(
            pos_args=self._parse_field_list(node.params),
            return_type=(
                [
                    self.eval_single_node(result_field.type)
                    for result_field in node.results.field_list
                ]
                if (node.results is not None and node.results.field_list is not None)
                else []
            ),
        )

    def _handle_basic_lit(self, node: asty_nodes.BasicLitNode) -> uast_nodes.Literal:
        return uast_nodes.Literal(node.value, node.kind)

    def _handle_composite_lit(self, node: asty_nodes.CompositeLitNode):
        match node.type:
            case asty_nodes.ArrayTypeNode():
                return uast_nodes.ArrayInitializer(
                    [self.eval_single_node(elem) for elem in node.elts]
                )
            case _:
                type_ast = self.eval_single_node_allowing_none(node.type)
                return uast_nodes.StructInitializer(
                    type_ast, [self.eval_single_node(elem) for elem in node.elts]
                )
                # raise NotImplementedError(node)

    def _handle_func_lit(self, node: asty_nodes.FuncLitNode) -> uast_nodes.MethodDecl:

        return uast_nodes.MethodDecl(
            None,
            type=self.eval_single_node(node.type),
            body=self.eval_single_node(node.body),
        )

    def _handle_ident(self, node: asty_nodes.IdentNode) -> uast_nodes.Name:
        """
        `ident` indicates identifier
        """
        return uast_nodes.Name(node.name)

    def _handle_binary_expr(
        self, node: asty_nodes.BinaryExprNode
    ) -> uast_nodes.BinaryExpr:
        op = node.op
        lhs_ast, rhs_ast = self.eval_single_node(node.x), self.eval_single_node(node.y)
        return uast_nodes.BinaryExpr(op, lhs_ast, rhs_ast)

    def _handle_unary_expr(
        self, node: asty_nodes.UnaryExprNode
    ) -> uast_nodes.UnaryExpr:
        op = node.op
        operand = self.eval_single_node(node.x)
        if node.position is not None:
            op_before_expr = node.position.offset < node.x.position.offset
        else:
            op_before_expr = True
        # if node.op_pos is None:
        return uast_nodes.UnaryExpr(op, operand, op_before_expr)

    def _handle_key_value_expr(
        self, node: asty_nodes.KeyValueExprNode
    ) -> uast_nodes.KeyValuePair:
        return uast_nodes.KeyValuePair(
            self.eval_single_node(node.key), self.eval_single_node(node.value)
        )

    def _handle_selector_expr(
        self, node: asty_nodes.SelectorExprNode
    ) -> uast_nodes.FieldAccessExpr:
        return uast_nodes.FieldAccessExpr(
            self.eval_single_node(node.x), self.eval_single_node(node.sel)
        )

    def _handle_call_expr(self, node: asty_nodes.CallExprNode) -> uast_nodes.CallExpr:
        return uast_nodes.CallExpr(
            name=self.eval_single_node(node.fun),
            arguments=self.eval_node_list(node.args) if node.args is not None else [],
        )

    def _handle_if_stmt(self, node: asty_nodes.IfStmtNode) -> uast_nodes.IfThenElseStmt:
        cond_ast = self.eval_single_node(node.cond)
        if_true_ast = self.eval_single_node(node.body)
        or_else_ast = (
            self.eval_single_node(node.else_pos) if node.else_pos is not None else None
        )
        return uast_nodes.IfThenElseStmt(
            cond_ast,
            if_true=if_true_ast,
            if_false=or_else_ast,
        )

    def _handle_case_clause(
        self, node: asty_nodes.CaseClauseNode
    ) -> uast_nodes.SwitchCase:
        def eval_case_body(expr_list: list) -> list[uast_nodes.SourceElement]:
            nodes = self.eval_node_list(expr_list)
            if (
                len(nodes) > 0
                and not isinstance(
                    nodes[-1],
                    (
                        uast_nodes.BreakStmt,
                        uast_nodes.ReturnStmt,
                        uast_nodes.ContinueStmt,
                    ),
                )
            ) or len(nodes) == 0:
                # Golang adds a break in every cases by default
                nodes.append(uast_nodes.BreakStmt())
            # print(nodes)
            return nodes

        return uast_nodes.SwitchCase(
            (
                self.eval_node_list(node.expr_list)
                if node.expr_list is not None
                else [uast_nodes.DefaultStmt()]
            ),
            uast_nodes.BlockStmt(eval_case_body(node.body)),
        )

    def _handle_switch_stmt(
        self, node: asty_nodes.SwitchStmtNode
    ) -> uast_nodes.SwitchStmt:
        predicate_ast = self.eval_single_node(node.tag)
        assert node.body is not None
        # case_clause: asty_nodes.CaseClauseNode
        cases_ast = [
            self.eval_single_node(case_clause) for case_clause in node.body.stmt_list
        ]
        return uast_nodes.SwitchStmt(
            predicate_ast,
            cases_ast,
        )

    def _handle_expr_stmt(self, node: asty_nodes.ExprStmtNode) -> uast_nodes.ExprStmt:
        return uast_nodes.ExprStmt(self.eval_single_node(node.x))

    def _handle_assign_stmt(
        self, node: asty_nodes.AssignStmtNode
    ) -> uast_nodes.Assignment:
        assert node.tok is not None
        return uast_nodes.Assignment(
            node.tok,
            lhs=self.eval_node_list(node.lhs),
            rhs=self.eval_node_list(node.rhs),
        )

    def _parse_value_spec(
        self, node: asty_nodes.ValueSpecNode
    ) -> list[uast_nodes.VarDecl]:
        return [
            uast_nodes.VarDecl(
                name.name,
                type=(
                    self.eval_single_node(node.type) if node.type is not None else None
                ),
                initializer=(
                    self.eval_single_node(node.values[i])
                    if node.values is not None
                    else None
                ),
            )
            for i, name in enumerate(node.names)
        ]

    def _handle_gen_decl(self, node: asty_nodes.GenDeclNode):
        match node:
            case asty_nodes.GenDeclNode(tok="var"):
                value_specs = []
                for spec in node.specs:
                    value_specs.extend(self._parse_value_spec(spec))
                return uast_nodes.CompoundDecl(value_specs)
            case asty_nodes.GenDeclNode(tok="type"):
                type_specs = []
                spec: asty_nodes.TypeSpecNode
                for spec in node.specs:
                    name_ast = self.eval_single_node(spec.name)
                    type_ast: uast_nodes.StructDecl = self.eval_single_node(spec.type)
                    type_ast.name = name_ast
                    type_specs.append(type_ast)
                return uast_nodes.CompoundDecl(type_specs)
            case _:
                raise NotImplementedError(node)

    def _handle_decl_stmt(self, node: asty_nodes.DeclStmtNode) -> uast_nodes.VarDecl:
        return self.eval_single_node(node.decl)

    def _handle_block_stmt(
        self, node: asty_nodes.BlockStmtNode
    ) -> uast_nodes.BlockStmt:
        return uast_nodes.BlockStmt(
            [self.eval_single_node(stmt) for stmt in node.stmt_list]
        )

    def _handle_func_decl(self, node: asty_nodes.FuncDeclNode):
        return uast_nodes.MethodDecl(
            self.eval_single_node(node.name),
            type=self.eval_single_node(node.type),
            body=self.eval_single_node(node.body),
        )

    def _parse_gen_decl(
        self, node: asty_nodes.GenDeclNode
    ) -> list[uast_nodes.ImportDecl]:
        data = []
        for spec in node.specs:
            if isinstance(spec, asty_nodes.ImportSpecNode):
                data.append(
                    uast_nodes.ImportDecl(
                        name=spec.name, path=self.eval_single_node(spec.path)
                    )
                )
            else:
                raise NotImplementedError
        return data
        # return uast_nodes.ImportDecl()

    def _handle_file(self, node: asty_nodes.FileNode) -> uast_nodes.CompilationUnit:
        data = []
        for decl in node.decls:
            match decl.node_type:
                case NodeType.GEN_DECL:
                    pass
                case _:
                    data.append(self.eval_single_node(decl))
        return uast_nodes.CompilationUnit(data)

    def eval(self, node: asty_nodes.BaseNode) -> uast_nodes.SourceElement:
        return self.eval_single_node(node)

    def eval_single_node(self, node: asty_nodes.BaseNode) -> uast_nodes.SourceElement:
        try:
            ret = self._handlers_map[node.node_type](node)

            assert isinstance(ret, uast_nodes.SourceElement)
            if node.position is not None:
                ret.location = (node.position.line, node.position.column)
            return ret
        except KeyError as e:
            print(node)
            raise e

    def eval_single_node_allowing_none(
        self, node: Optional[asty_nodes.BaseNode]
    ) -> Optional[uast_nodes.SourceElement]:
        """
        Evaluate node. If node is None, return None.
        """
        return self.eval_single_node(node) if node is not None else None

    def eval_node_list(
        self, nodes: Iterable[asty_nodes.BaseNode]
    ) -> list[uast_nodes.SourceElement]:
        return [self.eval_single_node(node) for node in nodes]
